// implement fork from user space

#include <inc/string.h>
#include <inc/lib.h>

// PTE_COW marks copy-on-write page table entries.
// It is one of the bits explicitly allocated to user processes (PTE_AVAIL).
#define PTE_COW		0x800

//
// Custom page fault handler - if faulting page is copy-on-write,
// map in our own private writable copy.
//
inline void print_utf(struct UTrapframe *utf)
{
	cprintf("fault_va %x\nerr %x\neip %x\neflags %x\nesp %x\n",
		utf->utf_fault_va, utf->utf_err, utf->utf_eip, utf->utf_eflags, utf->utf_esp);
}
static void
pgfault(struct UTrapframe *utf)
{
	void *addr = (void *) utf->utf_fault_va;
	uint32_t err = utf->utf_err;
	int r;

	// Check that the faulting access was (1) a write, and (2) to a
	// copy-on-write page.  If not, panic.
	// Hint:
	//   Use the read-only page table mappings at vpt
	//   (see <inc/memlayout.h>).

	// LAB 4: Your code here.
	if (!(err & FEC_WR)) {
		print_utf(utf);
		panic("pgfault: not a write access %x %x %x %x\n", sys_getenvid(), addr, err, utf->utf_eip);
	}
	if (!(vpt[PPN(addr)] & PTE_COW)) {
		print_utf(utf);
		panic("pgfault: not a copy-on-write page %x %x %x %x\n", sys_getenvid(), addr, err, utf->utf_eip);
	}


	// Allocate a new page, map it at a temporary location (PFTEMP),
	// copy the data from the old page to the new page, then move the new
	// page to the old page's address.
	// Hint:
	//   You should make three system calls.
	//   No need to explicitly delete the old page's mapping.

	// LAB 4: Your code here.
	// ROUND DOWN ADDR!!!
	addr = (void *)ROUNDDOWN(addr, PGSIZE);

	r = sys_page_alloc(0, (void *)PFTEMP, PTE_U | PTE_W | PTE_P);
	if (r < 0)
		panic("pgfault: page_alloc: %e\n", r);
	memmove((void *)PFTEMP, addr, PGSIZE);
	struct SyscallP arg[2];
	arg[0].syscallno = SYS_page_map;
	arg[0].a1 = 0;
	arg[0].a2 = (uint32_t)PFTEMP;
	arg[0].a3 = 0;
	arg[0].a4 = (uint32_t)addr;
	arg[0].a5 = PTE_U | PTE_W | PTE_P;
	arg[1].syscallno = SYS_page_unmap;
	arg[1].a1 = 0;
	arg[1].a2 = (uint32_t)PFTEMP;
	sys_batch(arg, arg + 1, 0, 0, 0);
	if (arg[0].ret < 0 || arg[1].ret < 0)
		panic("pgfault: mapping: %e\n", arg[0].ret < 0 ?
		      arg[0].ret : arg[1].ret);
	/*
	r = sys_page_map(0, (void *)PFTEMP, 0, addr, PTE_U | PTE_W | PTE_P);
	if (r < 0)
		panic("pgfault: page_map: %e\n", r);

	r = sys_page_unmap(0, (void *)PFTEMP);
	if (r < 0)
		panic("pgfault: page_unmap: %e\n", r);
	*/
//	cprintf("pgfault handled %x @ %x\n", sys_getenvid(), addr);
}

//
// Map our virtual page pn (address pn*PGSIZE) into the target envid
// at the same virtual address.  If the page is writable or copy-on-write,
// the new mapping must be created copy-on-write, and then our mapping must be
// marked copy-on-write as well.  (Exercise: Why mark ours copy-on-write again
// if it was already copy-on-write?)
//
// Returns: 0 on success, < 0 on error.
// It is also OK to panic on error.
//
static int
duppage(envid_t envid, unsigned pn)
{
	int r;
	void *addr = (void *)(pn << PTXSHIFT);
	pte_t pte = vpt[pn];

	// LAB 4: Your code here.
	if (pte & PTE_SHARE) {
		r = sys_page_map(0, addr, envid, addr, pte & PTE_USER);
		if (r < 0)
			panic("duppage: cow: page_map: %x %x %e\n",
			      envid, pn, r);
	} else if (pte & (PTE_W | PTE_COW)) {
		r = sys_page_map(0, addr, envid, addr, PTE_P | PTE_U | PTE_COW);
		if (r < 0)
			panic("duppage: cow: page_map: %x %x %e\n", envid, pn, r);

		r = sys_page_map(0, addr, 0, addr, PTE_P | PTE_U | PTE_COW);
		if (r < 0)
			panic("duppage: cow: page_map: %x %x %e\n", envid, pn, r);
	} else {
		r = sys_page_map(0, addr, envid, addr, PTE_P | PTE_U);
		if (r < 0)
			panic("duppage: ro: page_map: %x %x %e\n", envid, pn, r);
	}
	return 0;
}

//
// User-level fork with copy-on-write.
// Set up our page fault handler appropriately.
// Create a child.
// Copy our address space and page fault handler setup to the child.
// Then mark the child as runnable and return.
//
// Returns: child's envid to the parent, 0 to the child, < 0 on error.
// It is also OK to panic on error.
//
// Hint:
//   Use vpd, vpt, and duppage.
//   Remember to fix "env" and the user exception stack in the child process.
//   Neither user exception stack should ever be marked copy-on-write,
//   so you must allocate a new page for the child's user exception stack.
//
envid_t
fork(void)
{
	// LAB 4: Your code here.
	int r;
	envid_t child = -1;

	set_pgfault_handler(pgfault);

	r = sys_exofork();
	if (r < 0)
		panic("fork: exofork: %e\n", r);
	else if (r == 0) {
		env = envs + ENVX(sys_getenvid());
		return 0;
	} else
		child = r;

	int i, j;
	for (i = 0; i <= PDX(USTACKTOP - PGSIZE); ++i)
		if (vpd[i] & PTE_P)
			for (j = PPN((i << 22)); j < PPN((i << 22) + PTSIZE) && j < PPN(UXSTACKTOP - PGSIZE); ++j)
				if (vpt[j] & PTE_P)
					duppage(child, j);

	struct SyscallP syscall_arg[3];
	/*
	r = sys_page_alloc(child, (void *)(UXSTACKTOP - PGSIZE), PTE_P | PTE_U | PTE_W);
	if (r < 0)
		panic("fork: page_alloc: %e\n", r);
	*/
	syscall_arg[0].syscallno = SYS_page_alloc;
	syscall_arg[0].a1 = child;
	syscall_arg[0].a2 = UXSTACKTOP - PGSIZE;
	syscall_arg[0].a3 = PTE_P | PTE_U | PTE_W;

	/*
	r = sys_env_set_pgfault_upcall(child, env->env_pgfault_upcall);
	if (r < 0)
		panic("fork: pgfault_upcall: %e\n", r);
	*/
	syscall_arg[1].syscallno = SYS_env_set_pgfault_upcall;
	syscall_arg[1].a1 = child;
	syscall_arg[1].a2 = (uint32_t)env->env_pgfault_upcall;

	/*
	r = sys_env_set_status(child, ENV_RUNNABLE);
	if (r < 0)
		panic("fork: set_status: %e\n", r);
	*/
	syscall_arg[2].syscallno = SYS_env_set_status;
	syscall_arg[2].a1 = child;
	syscall_arg[2].a2 = ENV_RUNNABLE;

	sys_batch(syscall_arg, syscall_arg + 1, syscall_arg + 2, 0, 0);
	if (syscall_arg[0].ret < 0)
		panic("fork: page_alloc: %e\n", syscall_arg[0].ret);
	if (syscall_arg[1].ret < 0)
		panic("fork: pgfault_upcall: %e\n", syscall_arg[1].ret);
	if (syscall_arg[2].ret < 0)
		panic("fork: set_status: %e\n", syscall_arg[2].ret);

	return child;
}

// Share most pages unless it's a stack page
static void
sduppage(envid_t env, uint32_t pn)
{
	int r;
	void *addr = (void *)(pn << PTXSHIFT);
	pte_t pte = vpt[pn];
	int perm = pte & 0xe07;
	if ((uint32_t)addr < USTACKTOP
	    && (uint32_t)addr >= ROUNDDOWN((uint32_t)&perm, PGSIZE)) {
		//		cprintf("%x is stack\n", addr);
		r = sys_page_alloc(0, (void *)PFTEMP, perm);
		if (r < 0)
			panic("sduppage: stack: page_alloc: %e\n", r);
		memmove((void *)PFTEMP, addr, PGSIZE);
		/*
		r = sys_page_map(0, (void *)PFTEMP, env, addr, perm);
		if (r < 0)
			panic("sduppage: stack: page_map: %e\n", r);
		r = sys_page_unmap(0, (void *)PFTEMP);
		if (r < 0)
			panic("sduppage: stack: page_unmap: %e\n", r);
		*/
		struct SyscallP arg[2];
		arg[0].syscallno = SYS_page_map;
		arg[0].a1 = 0;
		arg[0].a2 = (uint32_t)PFTEMP;
		arg[0].a3 = env;
		arg[0].a4 = (uint32_t)addr;
		arg[0].a5 = (uint32_t)perm;
		arg[1].syscallno = SYS_page_unmap;
		arg[1].a1 = 0;
		arg[1].a2 = (uint32_t)PFTEMP;
		sys_batch(arg, arg + 1, 0, 0, 0);
		if (arg[0].ret < 0 || arg[1].ret < 0)
			panic("sduppage: stack: %e\n", arg[0].ret < 0 ?
			      arg[0].ret : arg[1].ret);
	} else {
		//		cprintf("%x is not stack\n", addr);
		r = sys_page_map(0, addr, env, addr, perm);
		if (r < 0)
			panic("sduppage: nonstack: page_map: %e\n", r);
	}
}
// Challenge!
int
sfork(void)
{
	int r;
	envid_t child = -1;

	// Create a new env
	r = sys_exofork();
	if (r < 0)
		panic("sfork: exofork: %e\n", r);
	else if (r == 0) {
		// we are not setting up env
		return 0;
	} else
		child = r;

	// Duplicate the mapping
	int i, j;
	for (i = 0; i <= PDX(USTACKTOP - PGSIZE); ++i)
		if (vpd[i] & PTE_P)
			for (j = PPN((i << 22)); j < PPN((i << 22) + PTSIZE) && j < PPN(UXSTACKTOP - PGSIZE); ++j)
				if (vpt[j] & PTE_P)
					sduppage(child, j);

	/*
	// Use the same pgfault_upcall as the parent env
	r = sys_env_set_pgfault_upcall(child, env->env_pgfault_upcall);
	if (r < 0)
		panic("sfork: set_pgfault_upcall: %e\n", r);

	// Mark it runnable
	r = sys_env_set_status(child, ENV_RUNNABLE);
	if (r < 0)
		panic("sfork: set_status: %e\n", r);
	*/
	struct SyscallP syscall_arg[2];
	syscall_arg[0].syscallno = SYS_env_set_pgfault_upcall;
	syscall_arg[0].a1 = child;
	syscall_arg[0].a2 = (uint32_t)env->env_pgfault_upcall;

	syscall_arg[1].syscallno = SYS_env_set_status;
	syscall_arg[1].a1 = child;
	syscall_arg[1].a2 = ENV_RUNNABLE;

	sys_batch(syscall_arg, syscall_arg + 1, 0, 0, 0);
	if (syscall_arg[0].ret < 0)
		panic("sfork: set_pgfault_upcall: %e\n", syscall_arg[0].ret);
	if (syscall_arg[1].ret < 0)
		panic("sfork: set_env_status: %e\n", syscall_arg[1].ret);

	return child;
}
